#include "ModularRobot.h"

#include "MyMath.h"

#include "HyperNEAT.h"
#include "RL_PoWER.h"
#include "SplineNeat.h"

#include <boost/filesystem.hpp>
#include <boost/property_tree/exceptions.hpp>
#include <boost/property_tree/json_parser.hpp>
#include <boost/property_tree/ptree.hpp>

#include <stdexcept>
#include <ctime>

const std::string ModularRobot::GPS_NAME = "GPS";
const std::string ModularRobot::EMITTER_NAME = "Emitter";
const std::string ModularRobot::RECEIVER_NAME = "Receiver";

const std::string ModularRobot::BUFFER_NAME = "Buffer";
const std::string ModularRobot::RESULTS_DIR = "Results";

const int ModularRobot::GPS_LOG_PERIDO = 50;

const int ModularRobot::SUPERVISOR_CHANNEL = 1024;

ModularRobot::ModularRobot()
:
_parameters(_init_parameters(getControllerArguments())),
_name(_parameters->get<std::string>("Robot.Name")),
_seed(0),
_r_index(_parameters->get<std::size_t>("Robot.Index")),
_r_index_root(_parameters->get<std::size_t>("Robot.Index_Root")),
organismSize(_parameters->get<std::size_t>("Robot.Modules_#")),
_m_index(_parameters->get<std::size_t>("Robot." + getName() + ".Index") - _r_index),
_m_type(_parameters->get<int>("Robot." + getName() + ".Type")),
_time_step(getBasicTimeStep()),
_time_start(0.0),
_time_offset(0.0),
_time_end(0.0),
_ev_type(_parameters->get<int>("Algorithm.Type")),
_ev_step(0),
_ev_steps_recovery(_parameters->get<std::size_t>("Algorithm.Recovery_Steps")),
_ev_steps_total(_parameters->get<std::size_t>("Algorithm.Evaluation_Steps") + _ev_steps_recovery),
totalEvaluations(_parameters->get<unsigned int>("Algorithm.Evaluations")),
_ev_angular_velocity(_parameters->get<double>("Algorithm.Angular_Velocity")),
_position_start(transforms::Vector_3::ZERO),
_position_end(transforms::Vector_3::ZERO),
numMotors(0),
motorRange(0),
_algorithm(0),
_gps(isRoot() ? _init_gps(_time_step) : 0),
_motors(_m_type ? _init_motors(_time_step) : 0)
{
    std::cout << "[" << getTime() << "] " << getName() << ": " << "Starting Controller, timestep: " << _time_step  << std::endl;
#ifdef DEBUG_CONTROLLER
    std::cout << "[" << getTime() << "] " << getName() << ": "
            << _parameters << " "
            << _name << " "
            << _seed << " "
            << _r_index << " "
            << _r_index_root << " "
            << _r_size << " "
            << _m_index << " "
            << _m_type << " "
            << _time_step << " "
            << _time_start << " "
            << _time_offset << " "
            << _time_end << " "
            << _ev_type << " "
            << _ev_step << " "
            << _ev_steps_recovery << " "
            << _ev_steps_total << " "
            << _position_start << " "
            << _position_end << " "
            << _algorithm << " "
            << _gps << " "
            << _emitter << " "
            << _receiver << " "
            << _motors << " "
            << std::endl;
#endif
    if (!isRoot()) {
        if (_parameters) {
            delete(_parameters);
        }
        _emitter = _init_emitter(static_cast<int> (_r_index_root));
        _receiver = _init_receiver(_time_step, static_cast<int> (_r_index_root+1));
        return;
    }else{
        _emitter = _init_emitter(static_cast<int> (_r_index_root+1));
        _receiver = _init_receiver(_time_step, static_cast<int> (_r_index_root));
    }

    bool flag = true;

    utils::Random rng;

    std::string parametersPath = _parameters->get<std::string>("Algorithm.Parameters");
    std::string save_path = _parameters->get<std::string>("Algorithm.Save");
    logDirectory = _parameters->get<std::string>("Algorithm.LogDir","");

    switch (_ev_type) {
        case ModularRobot::A_NEAT:

            while (flag) {
                flag = false;

                try {
                    _seed = rng.seed();
                    logDirectory = _init_directory(getProjectPath(), _name, "HyperNEAT", logDirectory, _seed);
                }

                catch (...) {
                    flag = true;
                }
            }

            if (save_path.empty()) {
                _algorithm = new HyperNEAT(
                                           _seed
                                           , parametersPath
                                           , logDirectory
                                           , _init_shape_size(_parameters->get_child("Shape"))
                                           , initialiseModuleMapping(_r_index, organismSize, _name, _parameters->get_child("Robot"))
                                           , numMotors
                                           , 1
                                        );
            }

            else {
                //_algorithm = new HyperNEAT(save_path);
            }

            break;

        case ModularRobot::A_POWER:

            while (flag) {
                flag = false;

                try {
                    _seed = rng.seed();
                    logDirectory = _init_directory(getProjectPath(), _name, "RL_PoWER", logDirectory, _seed);
                }

                catch (...) {
                    flag = true;
                }
            }

            /*
             * The Webots time step is being converted from milliseconds to
             * seconds
             */

            if (save_path.empty()) {
                _algorithm = new RL_PoWER(
                                        _seed,
                                        "",
                                        logDirectory,
                                        _time_step * 1e-3,
                                        _ev_angular_velocity,
                                        totalEvaluations,
                                        organismSize * numMotors
                                        );
            }

            else {
                _algorithm = new RL_PoWER(save_path,
                                        _time_step * 1e-3,
                                        _ev_angular_velocity,
                                        organismSize * numMotors);
            }

            break;

        case ModularRobot::A_CPG:
            throw std::runtime_error("Algorithm Unsupported");
            break;

        case ModularRobot::A_SPLINENEAT:
            
            while (flag) {
                flag = false;
                
                try {
                    _seed = rng.seed();
                    logDirectory = _init_directory(getProjectPath(), _name, "SplineNEAT", logDirectory, _seed);
                }
                
                catch (...) {
                    flag = true;
                }
            }
            
            _algorithm = new SplineNeat(  _seed
                                        , parametersPath
                                        , logDirectory
                                        , initialiseModuleMapping(_r_index, organismSize, _name, _parameters->get_child("Robot"))
                                        , numMotors
                                        , _time_step * 1e-3
                                        , _ev_angular_velocity
                                        , (unsigned int)organismSize);
            break;
        default:
            throw std::runtime_error("Algorithm Unknown");
            break;
    }
    
    std::cout << "Algorithm: " << &_algorithm << std::endl;

    /* Avoid wasting memory */
    if (_parameters) {
        delete(_parameters);
    }
}

ModularRobot::~ModularRobot()
{
    if (_algorithm) {
        delete(_algorithm);
    }
}

void ModularRobot::run()
{
    if (step(_time_step) == -1) {
        return;
    }

    if (isRoot()) {
#ifdef DEBUG_CONTROLLER
        std::cout << "[" << getTime() << "] " << getName() << ": We are root initialize things" << std::endl;
#endif
        
        bool flag = false;
        std::pair<double, std::string> fitness(0.0, "");
        
        boost::filesystem::path dirpath(logDirectory);
        
        std::cout << "Opening fitness log file: " << dirpath / "fitness.log" << std::endl;
        std::ofstream fitnessLog((dirpath / "fitness.log").c_str(), std::ofstream::app);
        
        std::cout << "Opening location log file: " << dirpath / "fitness.log" << std::endl;
        gpsLog = std::ofstream((dirpath / "location.log").c_str(), std::ofstream::app);
        
        fitnessLog << "#Generation Evaluation Fitness 2eDisSq Distance Speed RLFitness  (Seed: " << _seed << ")" << std::endl;
        gpsLog << "#Evaluation x y z" << std::endl;
        
        doubledvector anglesIn;
        doubledvector anglesOut;
        dvector anglesTMinusOne;
        dvector anglesTPlusOne;
        // Initialise the motor vectors.
        // AnglesIn are sent to the controller
        // AnglesOut are sent to the other modules
        anglesIn = doubledvector(organismSize,std::vector<double>(numMotors,0.0));
        anglesOut = doubledvector(organismSize,std::vector<double>(numMotors,0.0));
        // AnglesTMinusOne and AnglesTPlusOne are for the root modules
        // because the other modules are 1 timestep behind in executing due to communication.
        anglesTMinusOne = dvector(numMotors,0.0);
        anglesTPlusOne = dvector(numMotors,0.0);

        if (step(_time_step) == -1) {
            return;
        }

        if (step(_time_step) == -1) {
            return;
        }

        for(size_t i=0;i<numMotors;i++){
            _set_motor_position(i,0.0);
        }

        _time_offset = getTime();
        double tMinusOne,tPlusOne;
        while (step(_time_step) != -1) {    
            anglesIn = receiveAngles();
#ifdef DEBUG_TIMING
            std::cout << "[" << getTime() << "] " << getName() << " using root angles of timestep: " << tMinusOne << std::endl;
#endif
            for(size_t i=0;i<numMotors;i++){
                anglesIn[_m_index][i] = anglesTMinusOne[i];
                anglesTMinusOne[i] = _get_motor_position(i);
                tMinusOne = getTime();
            }
            
            anglesOut = _compute_angles(anglesIn);
            
            sendAngles(anglesOut);
#ifdef DEBUG_TIMING
            std::cout << "[" << getTime() << "] " << getName() << " setting root angles of time: " << tPlusOne << std::endl;
#endif
            for(size_t i=0;i<numMotors;i++){
                _set_motor_position(i,anglesTPlusOne[i]);
                anglesTPlusOne[i] = anglesOut[_m_index][i];
                tPlusOne = getTime();
            }

            if (_ev_step < _ev_steps_total) {
                _ev_step += 1;

                if ((_ev_step - 1) == _ev_steps_recovery) {
                    _time_start = getTime();
                    _position_start = _get_gps();
                }
                
                if(_ev_step % GPS_LOG_PERIDO == 0){
                    logGPS();
                }

                continue;
            }

            _time_end = getTime();
            _position_end = _get_gps();

            fitness = _compute_fitness((_time_end - _time_start), (_position_end - _position_start));

            _algorithm->setEvaluationFitness(fitness.first);
            _algorithm->setEvaluationFitnessAlt(fitness.second);
            
            fitnessLog << _algorithm->getGeneration() << " " << _algorithm->getEvaluation() << " " << getRealFitness(fitness.first) << " " << fitness.second << std::endl;

            if ((!flag) && (!_algorithm->nextEvaluation())) {
                flag = true;
                _algorithm->save();

//                exportImage((boost::filesystem::path(logDirectory) / "Last_State.jpeg").c_str(), 100);
//                simulationRevert();
                
                _emitter->setChannel(SUPERVISOR_CHANNEL);
                std::string message = "simulationdone";
                std::cout << "Done with evaluations sending revert command to supervisor with message: " << message << " on channel " << SUPERVISOR_CHANNEL << std::endl;
                _emitter->send(message.c_str(),(int)message.length()+1);
                _emitter->setChannel(static_cast<int>(_r_index_root)+1);
            }

            _ev_step = 0;
        }
    } else {

        if (step(_time_step) == -1) {
            return;
        }

        // Initialise the motor vectors.
        // AnglesIn are received from the root
        // AnglesOut are sent to the root
        dvector anglesIn(numMotors,0.0);
        dvector anglesOut(numMotors,0.0);
        
        for(size_t i=0;i<numMotors;i++){
            _set_motor_position(i,0.0);
        }

        while (step(_time_step) != -1) {
            for(size_t i=0;i<numMotors;i++){
                anglesOut[i] = _get_motor_position(i);
            }
            
            sendAngles(_m_index,anglesOut);
            anglesIn = receiveAngles(_m_index);
#ifdef DEBUG_CONTROLLER
            std::cout << "[" << getTime() << "] " << getName() << " Received vector with: " << anglesIn.size() << " elements:" << std::endl;
            for(int j=0;j<MOTOR_SIZE;j++){
                std::cout << anglesIn[j] << ", ";
            }
            std::cout << std::endl;
#endif
            for(size_t i=0;i<numMotors;i++){
                _set_motor_position(i,anglesIn[i]);
            }
        }
    }
}

void ModularRobot::sendAngles(doubledvector anglesOut){
    if(anglesOut.size() > 0){
        using namespace boost::property_tree;
        ptree data;
        data.put("timestamp",getTime());
        for(int i=0;i<anglesOut.size();i++){
            ptree sub;
            for(int j=0;j<anglesOut[i].size();j++){
                  sub.put(std::to_string(j),anglesOut[i][j]);
            }
            data.add_child(std::to_string(i), sub);
        }
        
        std::ostringstream datastring;
        boost::property_tree::json_parser::write_json(datastring,data,false);
#ifdef DEBUG_CONTROLLER
        std::cout << "[" << getTime() << "] sendAngles " << getName()
        << " sending data: " << datastring.str() << std::endl;
#endif
        _emitter->send(datastring.str().c_str(),(int)datastring.str().size()+1);
    }
}

void ModularRobot::sendAngles(size_t index, dvector anglesOut){
    if(anglesOut.size() > 0){
        using namespace boost::property_tree;
        ptree data;
        data.put("index",std::to_string(index));
        data.put("timestamp",getTime());
        for(int j=0;j<anglesOut.size();j++){
            data.put(std::to_string(j),anglesOut[j]);
        }
        
        std::ostringstream datastring;
        boost::property_tree::json_parser::write_json(datastring,data,false);
#ifdef DEBUG_CONTROLLER
        std::cout << "[" << getTime() << "] sendAngles(index) " << getName()
        << " sending data: " << datastring.str() << std::endl;
#endif
        _emitter->send(datastring.str().c_str(),(int)datastring.str().size()+1);
    }
}

doubledvector ModularRobot::receiveAngles(){
    doubledvector result(organismSize,dvector(numMotors,0.0));

#ifdef DEBUG_CONTROLLER
    std::cout << "[" << getTime() << "] receiveAngles " << getName() << ": queue length: " << _receiver->getQueueLength()  << std::endl;
#endif
    while(_receiver->getQueueLength() > 0){
#ifdef DEBUG_CONTROLLER
        std::cout << "[" << getTime() << "] receiveAngles " << getName() << ": received data of size: " << _receiver->getDataSize()
        << " data: " << (char*)_receiver->getData() << std::endl;
#endif
        
        using namespace boost::property_tree;
        ptree data;
        std::istringstream stream((char*)_receiver->getData());
        boost::property_tree::json_parser::read_json(stream, data);
        
        int index = data.get<double>("index");
        for(int j=0;j<numMotors;j++){
            result[index][j] = data.get<double>(std::to_string(j));
        }
        
#ifdef DEBUG_CONTROLLER
        std::cout  std::cout << "[" << getTime() << "] " << getName() << " received vector with: " << result.size() << " elements" << std::endl;
#endif
        
#ifdef DEBUG_TIMING
        std::cout << "[" << getTime() << "] " << getName() << " received angles of timestep: " << data.get<double>("timestamp") << std::endl;
#endif
        
        _receiver->nextPacket();
    }
    
    return result;
}

dvector ModularRobot::receiveAngles(size_t index){
    dvector result(numMotors,0.0);
    
#ifdef DEBUG_CONTROLLER
    std::cout << "[" << getTime() << "] receiveAngles(index) " << getName() << ": queue length: " << _receiver->getQueueLength()  << std::endl;
#endif
    if(_receiver->getQueueLength()){
#ifdef DEBUG_CONTROLLER
        std::cout << "[" << getTime() << "] receiveAngles(index) " << getName() << ": received data of size: " << _receiver->getDataSize()
        << " data: " << (char*)_receiver->getData() << std::endl;
#endif
        
        using namespace boost::property_tree;
        std::istringstream stream((char*)_receiver->getData());
        ptree data;
        boost::property_tree::json_parser::read_json(stream, data);
        
        ptree sub = data.get_child(std::to_string(index));
        for(int j=0;j<numMotors;j++){
            result[j] = sub.get<double>(std::to_string(j));
        }
#ifdef DEBUG_CONTROLLER
        std::cout << "Received vector with: " << result.size() << " elements:" << std::endl;
        for(int j=0;j<MOTOR_SIZE;j++){
            std::cout << result[j] << ", ";
        }
        std::cout << std::endl;
#endif
        
#ifdef DEBUG_TIMING
        std::cout << "[" << getTime() << "] " << getName() << " got angles with timestamp: " << data.get<double>("timestamp") << std::endl;
#endif
        _receiver->nextPacket();
    }
#ifdef DEBUG_CONTROLLER
    else{
        std::cerr << "[" << getTime() << "] " << getName() << " did not receive any angle updates from root!" << std::endl;
    }
#endif
    
    return result;
}

transforms::Vector_3 ModularRobot::_get_gps()
{
    const double * values = _gps->getValues();
#ifdef DEBUG_CONTROLLER
    std::cout << "[" << getTime() << "] " << getName() << ": Get GPS Coordinates " << values[0] << " 0.0 " << values[2] << std::endl;
#endif
    return transforms::Vector_3(values[0], 0.0, values[2]);
}

double ModularRobot::_get_motor_position(size_t index)
{
    double value = (utils::MyMath::PI_HALF_INVERSE * _motors[index]->getPosition());
#ifdef DEBUG_CONTROLLER
    std::cout << "[" << getTime() << "] " << getName() << ": Get Servo Angle " << value << std::endl;
#endif
    return value;
}

void ModularRobot::_set_motor_position(size_t index, double value)
{
    if (value < -1.0) {
        value = -1.0;
    }
    if (value > 1.0) {
        value = 1.0;
    }

    value *= motorRange;

    _motors[index]->setPosition(value);

#ifdef DEBUG_CONTROLLER
    std::cout << "[" << getTime() << "] " << getName() << ": Set Servo Angle " << value << std::endl;
#endif
}


std::vector<std::vector<double> > ModularRobot::_compute_angles(const std::vector<std::vector<double> > & angles)
{
    size_t size = angles.size();
    std::vector<std::vector<double> > result(size,std::vector<double>(numMotors,0.0));

    _algorithm->reset();

    for (std::size_t index = 0; index < size; index++) {
        for(size_t j=0;j<numMotors;j++){
            _algorithm->setParameter((index*numMotors)+j, angles[index][j]);
        }
    }

    _algorithm->setParameter(organismSize*numMotors, std::sin(_ev_angular_velocity * (getTime() - _time_offset)));

    _algorithm->update();

    for (std::size_t index = 0; index < size; index++) {
        for(size_t j=0;j<numMotors;j++){
            result[index][j] = _algorithm->getParameter((index*numMotors)+j);
        }
    }

    return result;
}

std::pair<double, std::string> ModularRobot::_compute_fitness(double d_t, const transforms::Vector_3 & d_xyz)
{
    std::pair<double, std::string> result;

    switch (_ev_type) {
        case ModularRobot::A_NEAT:
        case ModularRobot::A_SPLINENEAT:
            result.first = std::pow(2, d_xyz.lenght_square());
            break;

        case ModularRobot::A_CPG:
            result.first = d_xyz.lenght();
            break;

        case ModularRobot::A_POWER:
            result.first = d_xyz.lenght() / d_t;
            break;

        default:
            throw std::runtime_error("Unknown Algorithm");
            break;
    }

    result.second = boost::lexical_cast<std::string>(std::pow(2, d_xyz.lenght_square()))
            + " " +
            boost::lexical_cast<std::string>(d_xyz.lenght())
            + " " +
            boost::lexical_cast<std::string>(d_xyz.lenght() / d_t)
            + " " +
            boost::lexical_cast<std::string>(std::pow(100 * (d_xyz.lenght() / d_t),6));

    return result;
}

double ModularRobot::getRealFitness(double fitness)
{
    double result;
    switch (_ev_type) {
        case ModularRobot::A_NEAT:
        case ModularRobot::A_SPLINENEAT:
        case ModularRobot::A_CPG:
            result = fitness;
            break;
            
        case ModularRobot::A_POWER:
            result = std::pow(100 * fitness,6);
            break;
            
        default:
            throw std::runtime_error("Unknown Algorithm");
            break;
    }
    
    return result;
}

boost::property_tree::ptree * ModularRobot::_init_parameters(const std::string & path)
{
    boost::property_tree::ptree * result = new boost::property_tree::ptree();

    boost::property_tree::read_json(path, (* result));

    return result;
}

GPS * ModularRobot::_init_gps(double time_step)
{
    GPS * gps = getGPS(ModularRobot::GPS_NAME);

    if (!gps) {
        throw std::runtime_error("Device Not Found");
    }

    gps->enable(time_step);
#ifdef DEBUG_CONTROLLER
    std::cout << "[" << getTime() << "] " << getName() << ": Init GPS" << std::endl;
#endif
    return gps;
}

Emitter * ModularRobot::_init_emitter(int channel)
{
    Emitter * emitter = getEmitter(ModularRobot::EMITTER_NAME);

    if (!emitter) {
        throw std::runtime_error("Device Not Found");
    }

    emitter->setChannel(channel);
#ifdef DEBUG_CONTROLLER
    std::cout << "[" << getTime() << "] " << getName() << ": Init Emitter" << std::endl;
#endif
    return emitter;
}

Receiver * ModularRobot::_init_receiver(double time_step, int channel)
{
    Receiver * receiver = getReceiver(ModularRobot::RECEIVER_NAME);

    if (!receiver) {
        throw std::runtime_error("Device Not Found");
    }

    receiver->enable(time_step);
    receiver->setChannel(channel);
#ifdef DEBUG_CONTROLLER
    std::cout << "[" << getTime() << "] " << getName() << ": Init Receiver" << std::endl;
#endif
    return receiver;
}

Motor** ModularRobot::_init_motors(double time_step)
{
    numMotors = _parameters->get<unsigned int>("Robot.Motor.Number");
    motorRange = _parameters->get<double>("Robot.Motor.Range");
    
    Motor** motor = new Motor*[numMotors];
    for(size_t i=0;i<numMotors;i++){
        std::string motorName("Robot.Motor."+std::to_string(i));
        
        motor[i] = getMotor(_parameters->get<std::string>(motorName));
        if (!motor[i]) {
            throw std::runtime_error("Device Not Found");
        }
        motor[i]->enablePosition(time_step);
    }
#ifdef DEBUG_CONTROLLER
    std::cout << "[" << getTime() << "] " << getName() << ": Init Servo" << std::endl;
#endif
    return motor;
}

std::string ModularRobot::_init_directory(const std::string & directory, const std::string & name, const std::string & algorithm, const std::string &logdir, unsigned int instance)
{
    boost::filesystem::path path(directory);

    path /= ModularRobot::RESULTS_DIR;
    path /= name;
    path /= algorithm;
    if(logdir.length() > 0){
        path /= logdir;
    }
    
    std::time_t rawtime;
    std::tm* timeinfo;
    char buffer [80];
    
    std::time(&rawtime);
    timeinfo = std::localtime(&rawtime);
    
    std::strftime(buffer,80,"%Y-%m-%d-%H.%M.%S",timeinfo);
    
    path /= buffer;
    
//    path /= boost::lexical_cast<std::string>(instance);

    if (boost::filesystem::exists(path) && boost::filesystem::is_directory(path)) {
        throw std::runtime_error("Directory Already Existing");
    }

    boost::filesystem::create_directories(path);

    return path.string();
}

transforms::Vector_3 ModularRobot::_init_shape_size(const boost::property_tree::ptree & root)
{
    transforms::Vector_3 temp;

    temp.import_from_ptree(root);

    return transforms::Vector_3(temp.z(), temp.x(), temp.y());
}

std::vector<transforms::Vector_3> ModularRobot::initialiseModuleMapping(std::size_t offset, std::size_t size, const std::string & name, const boost::property_tree::ptree & root)
{
    transforms::Vector_3 p_index;
    boost::property_tree::ptree parent = root, child;

    std::vector<transforms::Vector_3> result(size);

    for (std::size_t index = 0; index < size; index++) {
        /* Get Module Property Tree */
        child = parent.get_child(name + ":" + boost::lexical_cast<std::string>(index + offset));

        /* Get Module 3D Index */
        p_index.import_from_ptree(child.get_child("Index_3D"));

        result[index] = transforms::Vector_3(p_index.z(), p_index.x(), p_index.y());

        std::cout << index << " -> " << result[index] << std::endl;
    }
#ifdef DEBUG_CONTROLLER
    std::cout << "[" << getTime() << "] " << getName() << ": Init Servo Mapping" << std::endl;
#endif
    return result;
}

void ModularRobot::logGPS(){
    const double * values = _gps->getValues();
    gpsLog << _algorithm->getEvaluation() << " " << values[0] << " 0.0 " << values[2] << std::endl;
}

bool ModularRobot::isRoot() {
    return _m_type == M_ROOT;
}